class Rotor:
    def __init__(self, wiring, notch):
        self.wiring = wiring
        self.notch = notch
        self.position = 0

    def set_position(self, position):
        self.position = position

    def forward(self, letter):
        offset = (ord(letter) - ord('A') + self.position) % 26
        return chr((ord(self.wiring[offset]) - ord('A') - self.position) % 26 + ord('A'))

    def backward(self, letter):
        offset = (ord(letter) - ord('A') + self.position) % 26
        return chr((self.wiring.index(chr(offset + ord('A'))) - self.position) % 26 + ord('A'))

    def rotate(self):
        self.position = (self.position + 1) % 26
        return self.position == self.notch

class Reflector:
    def __init__(self, wiring):
        self.wiring = wiring

    def reflect(self, letter):
        return self.wiring[ord(letter) - ord('A')]

class Plugboard:
    def __init__(self, pairs):
        self.mapping = {}
        for pair in pairs:
            a, b = pair
            self.mapping[a] = b
            self.mapping[b] = a

    def process(self, letter):
        return self.mapping.get(letter, letter)

class EnigmaMachine:
    def __init__(self, rotors, reflector, plugboard):
        self.rotors = rotors
        self.reflector = reflector
        self.plugboard = plugboard

    def set_rotors(self, positions):
        for rotor, position in zip(self.rotors, positions):
            rotor.set_position(position)

    def encrypt(self, letter):
        letter = self.plugboard.process(letter)
        for rotor in self.rotors:
            letter = rotor.forward(letter)
        letter = self.reflector.reflect(letter)
        for rotor in reversed(self.rotors):
            letter = rotor.backward(letter)
        letter = self.plugboard.process(letter)
        self.rotate_rotors()
        return letter

    def decrypt(self, letter):
        return self.encrypt(letter)

    def rotate_rotors(self):
        if self.rotors[0].rotate():
            if len(self.rotors) > 1 and self.rotors[1].rotate():
                if len(self.rotors) > 2:
                    self.rotors[2].rotate()

# 定义转轮和反射器的配置
rotor_I = Rotor("EKMFLGDQVZNTOWYHXUSPAIBRCJ", 17)
rotor_II = Rotor("AJDKSIRUXBLHWTMCQGZNPYFVOE", 5)
rotor_III = Rotor("BDFHJLCPRTXVZNYEIWGAKMUSQO", 22)
reflector_B = Reflector("YRUHQSLDPXNGOKMIEBFZCWVJAT")
plugboard_pairs = [("A", "B"), ("C", "D")]

# 创建Enigma机实例
enigma = EnigmaMachine([rotor_I, rotor_II, rotor_III], reflector_B, Plugboard(plugboard_pairs))

# 设置转轮初始位置
initial_positions = [0, 0, 0]
enigma.set_rotors(initial_positions)

# 测试解密函数
ciphertext = "KQJYQ"   #原始加密文本
decrypted_text = ''.join(enigma.decrypt(char) for char in ciphertext)
print("原始加密文本:", ciphertext)
print("解密后的明文:", decrypted_text)
